#include <iostream>
#include <hip/hip_runtime.h>
#include <unistd.h>
#include <sys/wait.h>
#include <vector>

#define N 10000
#define BLOCKS 1024
#define THREADS 256
__global__ void add(int *a, int *b, int *c,int n) {
    int id = blockIdx.x * blockDim.x + threadIdx.x;
    if(id<n)
    c[id] = a[id] + b[id];
}

__global__ void sub(int *a, int *b, int *c, int n) {
    int id = blockIdx.x * blockDim.x + threadIdx.x;
    if(id<n)
    c[id] = a[id] - b[id];
}

void LaunchAdd(){
    int *a, *b, *c;
    hipMalloc(&a, N * sizeof(int));
    hipMalloc(&b, N * sizeof(int));
    hipMalloc(&c, N * sizeof(int));
    int *h_a,*h_b,*h_c;
    h_a = new int[N];
    h_b = new int[N];
    h_c = new int[N];
    for(int i = 0; i < N; ++i){
        h_a[i] = i;
        h_b[i] = i + 1;
    }
    hipMemcpy(a, h_a, N * sizeof(int), hipMemcpyHostToDevice);
    hipMemcpy(b, h_b, N * sizeof(int), hipMemcpyHostToDevice);
    hipLaunchKernelGGL(add, dim3(1), dim3(THREADS), 0, 0, a, b, c,N);
    hipDeviceSynchronize();
    hipMemcpy(h_c, c, N * sizeof(int), hipMemcpyDeviceToHost);
    delete[] h_a;
    delete[] h_b;
    delete[] h_c;
    hipFree(a);
    hipFree(b);
    hipFree(c);
}

void LaunchSub(){
    int *a, *b, *c;
    hipMalloc(&a, N * sizeof(int));
    hipMalloc(&b, N * sizeof(int));
    hipMalloc(&c, N * sizeof(int));
    int *h_a,*h_b,*h_c;
    h_a = new int[N];
    h_b = new int[N];
    h_c = new int[N];
    for(int i = 0; i < N; ++i){
        h_a[i] = i;
        h_b[i] = i + 1;
    }
    hipMemcpy(a, h_a, N * sizeof(int), hipMemcpyHostToDevice);
    hipMemcpy(b, h_b, N * sizeof(int), hipMemcpyHostToDevice);
    hipLaunchKernelGGL(sub, dim3(BLOCKS), dim3(THREADS), 0, 0, a, b, c,N);
    hipDeviceSynchronize();
    hipMemcpy(h_c, c, N * sizeof(int), hipMemcpyDeviceToHost);
    delete[] h_a;
    delete[] h_b;
    delete[] h_c;
    hipFree(a);
    hipFree(b);
    hipFree(c);
}

int main() {
    pid_t pid = fork();
    if(pid == 0){
        LaunchAdd();
    } else if (pid > 0){
        LaunchSub();
        wait(NULL);
    }
    return 0;
}